package leetcode;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;

/**
 * @program: datastructureandalogorithm
 * @description:
 * @author: hmx
 * @create: 2022-01-14 00:52
 **/
public class LeetCode373 {

    //优先队列
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        //存储两个数组的索引(i1, i2)
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[2]]);

        int m = nums1.length;
        int n = nums2.length;
        //把(0,0),(1,0),...(min(k, m),0)添加到队列中
        for (int i = 0; i < Math.min(k, m); ++i) {
            pq.offer(new int[]{i, 0});
        }

        List<List<Integer>> res = new ArrayList<>();
        while (k-- > 0 && !pq.isEmpty()) {
            int[] poll = pq.poll();
            res.add(Arrays.asList(nums1[poll[0]], nums2[poll[1]]));
            //(i1, i2 + 1)未越界就添加到优先队列中
            if (poll[1] + 1 < n) {
                pq.offer(new int[]{poll[0], poll[1] + 1});
            }
        }

        return res;
    }


}
